use crate::{
    ast::MapExpr,
    error::{CompilerError, Result},
    llvm_backend::LLVMBackend,
};
use inkwell::{
    types::{BasicType, BasicTypeEnum},
    values::BasicValueEnum,
    AddressSpace,
};

impl<'ctx> LLVMBackend<'ctx> {
    /// Generate map expression (map literal)
    /// Map structure: {keys: *K, values: *V, len: i64, cap: i64}
    pub(in crate::llvm_backend) fn generate_map_expression(
        &mut self,
        map_expr: &MapExpr,
    ) -> Result<BasicValueEnum<'ctx>> {
        // Get key and value types from the first entry or use default types
        if map_expr.entries.is_empty() {
            // Empty map - we need type information from the MapType
            // For now, return a simple empty map with i8* pointers
            return self.generate_empty_map();
        }

        // Get element types by evaluating the first entry
        let first_entry = &map_expr.entries[0];
        let first_key = self.generate_expression(&first_entry.key)?;
        let first_value = self.generate_expression(&first_entry.value)?;

        let key_type = first_key.get_type();
        let value_type = first_value.get_type();

        // Generate all keys and values
        let mut keys = vec![first_key];
        let mut values = vec![first_value];

        for entry in &map_expr.entries[1..] {
            let key = self.generate_expression(&entry.key)?;
            let value = self.generate_expression(&entry.value)?;

            // Type checking - ensure all keys have the same type
            if key.get_type() != key_type {
                return Err(CompilerError::internal(
                    "All map keys must have the same type",
                ));
            }
            if value.get_type() != value_type {
                return Err(CompilerError::internal(
                    "All map values must have the same type",
                ));
            }

            keys.push(key);
            values.push(value);
        }

        let length = self.context.i64_type().const_int(keys.len() as u64, false);
        let capacity = length; // For literals, cap == len

        // Calculate sizes
        let key_size = self.get_type_size_bytes(key_type)?;
        let value_size = self.get_type_size_bytes(value_type)?;

        // Allocate memory for keys array
        let key_array_size_bytes = self.builder.build_int_mul(
            capacity,
            self.context.i64_type().const_int(key_size, false),
            "key_array_size",
        )?;

        let malloc_func = self.function_map.get("__malloc").unwrap();
        let keys_malloc_call = self.builder.build_call(
            *malloc_func,
            &[key_array_size_bytes.into()],
            "map_keys_malloc",
        )?;
        let keys_ptr_i8 = keys_malloc_call
            .try_as_basic_value()
            .left()
            .unwrap()
            .into_pointer_value();

        // Cast i8* to K*
        let keys_ptr = self
            .builder
            .build_bitcast(
                keys_ptr_i8,
                key_type.ptr_type(AddressSpace::default()),
                "map_keys_ptr",
            )?
            .into_pointer_value();

        // Allocate memory for values array
        let value_array_size_bytes = self.builder.build_int_mul(
            capacity,
            self.context.i64_type().const_int(value_size, false),
            "value_array_size",
        )?;

        let values_malloc_call = self.builder.build_call(
            *malloc_func,
            &[value_array_size_bytes.into()],
            "map_values_malloc",
        )?;
        let values_ptr_i8 = values_malloc_call
            .try_as_basic_value()
            .left()
            .unwrap()
            .into_pointer_value();

        // Cast i8* to V*
        let values_ptr = self
            .builder
            .build_bitcast(
                values_ptr_i8,
                value_type.ptr_type(AddressSpace::default()),
                "map_values_ptr",
            )?
            .into_pointer_value();

        // Fill keys and values arrays
        for (i, (key, value)) in keys.iter().zip(values.iter()).enumerate() {
            let index = self.context.i64_type().const_int(i as u64, false);

            // Store key
            let key_ptr = unsafe {
                self.builder.build_gep(
                    key_type,
                    keys_ptr,
                    &[index],
                    &format!("map_key_{}_ptr", i),
                )?
            };
            self.builder.build_store(key_ptr, *key)?;

            // Store value
            let value_ptr = unsafe {
                self.builder.build_gep(
                    value_type,
                    values_ptr,
                    &[index],
                    &format!("map_value_{}_ptr", i),
                )?
            };
            self.builder.build_store(value_ptr, *value)?;
        }

        // Create map struct: { keys: *K, values: *V, len: i64, cap: i64 }
        let keys_ptr_type = key_type.ptr_type(AddressSpace::default());
        let values_ptr_type = value_type.ptr_type(AddressSpace::default());
        let len_type = self.context.i64_type();
        let cap_type = self.context.i64_type();

        let map_struct_type = self.context.struct_type(
            &[
                keys_ptr_type.into(),
                values_ptr_type.into(),
                len_type.into(),
                cap_type.into(),
            ],
            false,
        );

        // Allocate memory for the map struct
        let map_alloca = self.builder.build_alloca(map_struct_type, "map_alloca")?;

        // Store keys pointer (field 0)
        let keys_field_ptr = unsafe {
            self.builder.build_gep(
                map_struct_type,
                map_alloca,
                &[
                    self.context.i32_type().const_int(0, false),
                    self.context.i32_type().const_int(0, false),
                ],
                "map_keys_field",
            )?
        };
        self.builder.build_store(keys_field_ptr, keys_ptr)?;

        // Store values pointer (field 1)
        let values_field_ptr = unsafe {
            self.builder.build_gep(
                map_struct_type,
                map_alloca,
                &[
                    self.context.i32_type().const_int(0, false),
                    self.context.i32_type().const_int(1, false),
                ],
                "map_values_field",
            )?
        };
        self.builder.build_store(values_field_ptr, values_ptr)?;

        // Store len (field 2)
        let len_field_ptr = unsafe {
            self.builder.build_gep(
                map_struct_type,
                map_alloca,
                &[
                    self.context.i32_type().const_int(0, false),
                    self.context.i32_type().const_int(2, false),
                ],
                "map_len_field",
            )?
        };
        self.builder.build_store(len_field_ptr, length)?;

        // Store cap (field 3)
        let cap_field_ptr = unsafe {
            self.builder.build_gep(
                map_struct_type,
                map_alloca,
                &[
                    self.context.i32_type().const_int(0, false),
                    self.context.i32_type().const_int(3, false),
                ],
                "map_cap_field",
            )?
        };
        self.builder.build_store(cap_field_ptr, capacity)?;

        // Load and return the complete map struct
        let map_value = self
            .builder
            .build_load(map_struct_type, map_alloca, "map_value")?;
        Ok(map_value)
    }

    /// Generate empty map
    fn generate_empty_map(&mut self) -> Result<BasicValueEnum<'ctx>> {
        // Empty map with null pointers and zero length/capacity
        let null_ptr = self
            .context
            .i8_type()
            .ptr_type(AddressSpace::default())
            .const_null();
        let zero = self.context.i64_type().const_int(0, false);

        let map_struct_type = self.context.struct_type(
            &[
                self.context
                    .i8_type()
                    .ptr_type(AddressSpace::default())
                    .into(), // keys
                self.context
                    .i8_type()
                    .ptr_type(AddressSpace::default())
                    .into(), // values
                self.context.i64_type().into(), // len
                self.context.i64_type().into(), // cap
            ],
            false,
        );

        let map_alloca = self
            .builder
            .build_alloca(map_struct_type, "empty_map_alloca")?;

        // Store null keys pointer
        let keys_field_ptr = unsafe {
            self.builder.build_gep(
                map_struct_type,
                map_alloca,
                &[
                    self.context.i32_type().const_int(0, false),
                    self.context.i32_type().const_int(0, false),
                ],
                "empty_map_keys_field",
            )?
        };
        self.builder.build_store(keys_field_ptr, null_ptr)?;

        // Store null values pointer
        let values_field_ptr = unsafe {
            self.builder.build_gep(
                map_struct_type,
                map_alloca,
                &[
                    self.context.i32_type().const_int(0, false),
                    self.context.i32_type().const_int(1, false),
                ],
                "empty_map_values_field",
            )?
        };
        self.builder.build_store(values_field_ptr, null_ptr)?;

        // Store zero len
        let len_field_ptr = unsafe {
            self.builder.build_gep(
                map_struct_type,
                map_alloca,
                &[
                    self.context.i32_type().const_int(0, false),
                    self.context.i32_type().const_int(2, false),
                ],
                "empty_map_len_field",
            )?
        };
        self.builder.build_store(len_field_ptr, zero)?;

        // Store zero cap
        let cap_field_ptr = unsafe {
            self.builder.build_gep(
                map_struct_type,
                map_alloca,
                &[
                    self.context.i32_type().const_int(0, false),
                    self.context.i32_type().const_int(3, false),
                ],
                "empty_map_cap_field",
            )?
        };
        self.builder.build_store(cap_field_ptr, zero)?;

        let map_value = self
            .builder
            .build_load(map_struct_type, map_alloca, "empty_map_value")?;
        Ok(map_value)
    }

    /// Get size of a type in bytes
    fn get_type_size_bytes(&self, ty: BasicTypeEnum<'ctx>) -> Result<u64> {
        Ok(match ty {
            BasicTypeEnum::IntType(int_ty) => {
                match int_ty.get_bit_width() {
                    1 => 1,  // bool
                    8 => 1,  // i8
                    16 => 2, // i16
                    32 => 4, // i32
                    64 => 8, // i64
                    _ => 8,  // default
                }
            }
            BasicTypeEnum::FloatType(float_ty) => {
                // f32 or f64
                if float_ty == self.context.f32_type() {
                    4
                } else {
                    8
                }
            }
            BasicTypeEnum::PointerType(_) => 8, // 64-bit pointer
            _ => 8,                             // default
        })
    }
}
